#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
简化版污染类型数据导入脚本
直接操作数据库，避免复杂的验证逻辑
python scripts/import_pollution_types_simple.py "./data/污染类型数据-待初始化.xlsx" --clear-existing

"""

import sys
import os
from pathlib import Path
from typing import Dict, List
import pandas as pd
from sqlalchemy import create_engine, text

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from core.database import get_db_session
from core.logging_config import get_logger

logger = get_logger("import_pollution_types_simple")

class SimplePollutionTypesImporter:
    """简化版污染类型数据导入器"""

    def __init__(self):
        self.db_url = "sqlite:///./data/petition.db"
        self.engine = create_engine(self.db_url)

    def import_from_excel(self, excel_file_path: str, customer_id: str = "default", clear_existing: bool = False, db=None):
        """从Excel文件导入污染类型数据
        
        Args:
            excel_file_path: Excel文件路径
            customer_id: 客户ID，默认为"default"
            clear_existing: 是否清空现有数据，默认为False
            db: 可选的数据库会话，如果为None则自动创建
        """
        try:
            logger.info(f"开始从Excel文件导入污染类型数据: {excel_file_path}")

            # 检查文件是否存在
            if not os.path.exists(excel_file_path):
                raise FileNotFoundError(f"Excel文件不存在: {excel_file_path}")

            # 读取Excel文件
            df = pd.read_excel(excel_file_path, engine='openpyxl')
            logger.info(f"读取到 {len(df)} 行数据")

            # 验证必要的列 - 支持多种列名映射
            column_mappings = {
                '标签': ['标签', '名称', 'name', '标题'],
                '父级标签': ['父级标签', '父级', 'parent_code', 'parent'],
                '标签级别': ['标签级别', '标签层级', '层级', 'level', '㼶']
            }

            found_columns = {}
            for key, possible_names in column_mappings.items():
                found = False
                for name in possible_names:
                    if name in df.columns:
                        found_columns[key] = name
                        found = True
                        break
                if not found:
                    raise ValueError(f"缺少必要的列: {key} (期望列名: {possible_names})")

            logger.info(f"列名映射: {found_columns}")

            # 重命名列以统一处理
            df_renamed = df.rename(columns={
                found_columns['标签']: '标签',
                found_columns['父级标签']: '父级标签',
                found_columns['标签级别']: '标签级别'
            })

            # 如果需要清空现有数据
            if clear_existing:
                if db is not None:
                    # 使用传入的数据库会话清空数据
                    db.execute(text("DELETE FROM pollution_types WHERE customer_id = :customer_id"), 
                              {"customer_id": customer_id})
                else:
                    # 使用原有的逻辑清空数据
                    self.clear_existing_data(customer_id)
                logger.info("已清空现有污染类型数据")

            # 导入数据
            result = self._import_data(df_renamed, customer_id, db=db)

            # 如果使用传入的数据库会话，需要手动提交事务
            if db is not None:
                try:
                    db.commit()
                    logger.info("数据库事务已提交")
                except Exception as e:
                    logger.error(f"提交数据库事务失败: {e}")
                    db.rollback()
                    raise

            # 输出结果
            self._print_import_results(result)

            return result

        except Exception as e:
            logger.error(f"导入数据失败: {str(e)}", exc_info=True)
            raise

    def _import_data(self, df: pd.DataFrame, customer_id: str, db=None):
        """导入数据
        
        Args:
            df: DataFrame数据
            customer_id: 客户ID
            db: 可选的数据库会话，如果为None则自动创建
        """
        success_count = 0
        errors = []

        # 清洗数据
        df = self._clean_dataframe(df)

        # 构建标签映射
        tag_to_info = {}
        for _, row in df.iterrows():
            tag = str(row['标签']).strip()
            tag_to_info[tag] = row

        logger.info(f"构建了 {len(tag_to_info)} 个标签的映射")

        # 按照级别顺序导入，先导入级别1的，再导入级别2的，依此类推
        for level in sorted(df['标签级别'].unique()):
            logger.info(f"开始导入级别 {level} 的数据")

            level_data = df[df['标签级别'] == level]
            logger.info(f"级别 {level} 有 {len(level_data)} 条数据")

            for _, row in level_data.iterrows():
                try:
                    result = self._import_single_row(row, customer_id, tag_to_info, db=db)
                    if result:
                        success_count += 1
                        logger.info(f"成功导入: {row['标签']} (级别: {row['标签级别']})")
                    else:
                        errors.append(f"导入失败: {row['标签']}")
                except Exception as e:
                    error_msg = f"导入 {row['标签']} 失败: {str(e)}"
                    errors.append(error_msg)
                    logger.warning(error_msg, exc_info=True)

        return {
            "success_count": success_count,
            "error_count": len(errors),
            "errors": errors
        }

    def _import_single_row(self, row, customer_id: str, tag_to_info: Dict, db=None):
        """导入单行数据
        
        Args:
            row: Excel行数据
            customer_id: 客户ID
            tag_to_info: 标签信息映射
            db: 可选的数据库会话，如果为None则自动创建
        """
        try:
            # 提取基本信息
            tag = str(row['标签']).strip()
            parent_tag = str(row['父级标签']).strip() if pd.notna(row['父级标签']) and str(row['父级标签']) != 'nan' else None
            level = int(row['标签级别'])

            # 生成code
            code = self._generate_code(tag, parent_tag, level)

            # 查找父级code
            parent_code = None
            if parent_tag:
                parent_code = self._find_parent_code(parent_tag, tag_to_info)
            
            # 使用传入的数据库会话或创建新的会话
            if db is not None:
                # 使用传入的数据库会话
                existing = db.execute(text(
                    "SELECT id FROM pollution_types WHERE customer_id = :customer_id AND code = :code"
                ), {"customer_id": customer_id, "code": code}).fetchone()

                if existing:
                    logger.info(f"跳过已存在的记录: {tag} ({code})")
                    return True

                # 插入新记录
                insert_sql = """
                INSERT INTO pollution_types (
                    customer_id, code, name, parent_code, level,
                    tag_name, tag_category, usage_scope, exclusion_scope,
                    display_name, parent_tag_code, sort_order, description,
                    is_active, is_system, created_at, updated_at
                ) VALUES (
                    :customer_id, :code, :name, :parent_code, :level,
                    :tag_name, :tag_category, :usage_scope, :exclusion_scope,
                    :display_name, :parent_tag_code, :sort_order, :description,
                    :is_active, :is_system, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
                )
                """

                params = {
                    "customer_id": customer_id,
                    "code": code,
                    "name": tag,
                    "parent_code": parent_code,
                    "level": level,
                    "tag_name": tag,
                    "tag_category": self._get_cell_value(row, '标签分类'),
                    "usage_scope": self._get_cell_value(row, '使用范围'),
                    "exclusion_scope": self._get_cell_value(row, '排除范围'),
                    "display_name": tag,
                    "parent_tag_code": parent_code,
                    "sort_order": self._get_cell_value(row, '序号', 0),
                    "description": self._get_cell_value(row, '说明', ''),
                    "is_active": True,
                    "is_system": False
                }

                db.execute(text(insert_sql), params)
                db.commit()
                logger.info(f"成功插入: {tag} (级别: {level}, 父级: {parent_code})")
                return True
            else:
                # 使用原有的逻辑，自动创建数据库会话
                with get_db_session() as db:
                    # 检查是否已存在
                    existing = db.execute(text(
                        "SELECT id FROM pollution_types WHERE customer_id = :customer_id AND code = :code"
                    ), {"customer_id": customer_id, "code": code}).fetchone()

                    if existing:
                        logger.info(f"跳过已存在的记录: {tag} ({code})")
                        return True

                    # 插入新记录
                    insert_sql = """
                    INSERT INTO pollution_types (
                        customer_id, code, name, parent_code, level,
                        tag_name, tag_category, usage_scope, exclusion_scope,
                        display_name, parent_tag_code, sort_order, description,
                        is_active, is_system, created_at, updated_at
                    ) VALUES (
                        :customer_id, :code, :name, :parent_code, :level,
                        :tag_name, :tag_category, :usage_scope, :exclusion_scope,
                        :display_name, :parent_tag_code, :sort_order, :description,
                        :is_active, :is_system, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
                    )
                    """

                    params = {
                        "customer_id": customer_id,
                        "code": code,
                        "name": tag,
                        "parent_code": parent_code,
                        "level": level,
                        "tag_name": tag,
                        "tag_category": self._get_cell_value(row, '标签分类'),
                        "usage_scope": self._get_cell_value(row, '使用范围'),
                        "exclusion_scope": self._get_cell_value(row, '排除范围'),
                        "display_name": tag,
                        "parent_tag_code": parent_code,
                        "sort_order": self._get_cell_value(row, '序号', 0),
                        "description": self._get_cell_value(row, '说明', ''),
                        "is_active": True,
                        "is_system": False
                    }

                    db.execute(text(insert_sql), params)
                    db.commit()

                    logger.info(f"成功插入: {tag} (级别: {level}, 父级: {parent_code})")
                    return True

        except Exception as e:
            logger.error(f"导入单行数据失败: {str(e)}", exc_info=True)
            return False

    def _generate_code(self, tag: str, parent_tag: str, level: int) -> str:
        """生成code"""
        # 将中英文标签转换为英文代码
        def to_english_code(text):
            mapping = {
                '空气污染': 'AIR',
                '水污染': 'WATER',
                '土壤污染': 'SOIL',
                '噪声污染': 'NOISE',
                '工业废水': 'INDUSTRIAL-WATER',
                '工业废气': 'INDUSTRIAL-GAS',
                '工业固废': 'INDUSTRIAL-SOLID',
                '工业噪声': 'INDUSTRIAL-NOISE',
                '普通噪声': 'GENERAL-NOISE',
                '大气污染': 'AIR',
                '废水': 'WATER',
                '废气': 'GAS',
                '废渣': 'SOLID',
                '工业': 'INDUSTRIAL',
                '噪声': 'NOISE',
                '普通': 'GENERAL',
                '生活': 'LIVING',
                '交通': 'TRAFFIC',
                '施工': 'CONSTRUCTION',
                '农业': 'AGRICULTURE',
                '畜禽养殖': 'LIVESTOCK',
                '土壤': 'SOIL',
                '水体': 'WATER',
                '空气': 'AIR',
                '噪声': 'NOISE'
            }

            if text in mapping:
                return mapping[text]

            import re
            clean_text = re.sub(r'[^a-zA-Z0-9]', '', text)
            if clean_text:
                return clean_text.upper()
            else:
                return f"TYPE{hash(text) % 1000:03d}"

        base = to_english_code(tag)
        return f"{base}-{level}"

    def _find_parent_code(self, parent_tag: str, tag_to_info: Dict) -> str:
        """查找父级code"""
        # 标准化父级标签
        normalized_parent = parent_tag.replace('.0', '')

        # 从Excel数据中查找父级
        if normalized_parent in tag_to_info:
            parent_row = tag_to_info[normalized_parent]
            return self._generate_code(normalized_parent, None, int(parent_row['标签级别']))

        # 如果找不到，尝试直接查找
        with get_db_session() as db:
            result = db.execute(text(
                "SELECT code FROM pollution_types WHERE name = :name"
            ), {"name": normalized_parent}).fetchone()

            if result:
                return result[0]

        logger.warning(f"未找到父级标签: {parent_tag}")
        return None

    def _get_cell_value(self, row, column_name: str, default=None):
        """获取单元格值"""
        if column_name in row and pd.notna(row[column_name]):
            value = row[column_name]
            if str(value) != 'nan':
                return str(value).strip()
        return default

    def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        """清洗DataFrame"""
        # 删除空行
        df = df.dropna(how='all')

        # 重置索引
        df = df.reset_index(drop=True)

        # 清洗字符串列
        string_columns = ['标签', '父级标签', '标签分类', '使用范围', '排除范围', '说明']
        for col in string_columns:
            if col in df.columns:
                df[col] = df[col].apply(lambda x: str(x).strip() if pd.notna(x) and str(x) != 'nan' else None)

        # 确保标签级别是整数
        if '标签级别' in df.columns:
            df['标签级别'] = pd.to_numeric(df['标签级别'], errors='coerce').fillna(1).astype(int)

        return df

    def clear_existing_data(self, customer_id: str):
        """清空现有的污染类型数据"""
        with get_db_session() as db:
            # 删除所有污染类型数据
            deleted_types = db.execute(text(
                "DELETE FROM pollution_types WHERE customer_id = :customer_id"
            ), {"customer_id": customer_id}).rowcount

            db.commit()
            logger.info(f"已清空数据: {deleted_types}个污染类型")

    def _print_import_results(self, results: dict):
        """打印导入结果"""
        print("\n" + "="*60)
        print("污染类型数据导入结果")
        print("="*60)

        print(f"\n污染类型导入:")
        print(f"  成功: {results['success_count']} 条")
        print(f"  失败: {results['error_count']} 条")

        if results['errors']:
            print(f"\n错误详情:")
            for error in results['errors'][:5]:
                print(f"  - {error}")
            if len(results['errors']) > 5:
                print(f"  ... 还有 {len(results['errors']) - 5} 个错误")

        print(f"\n总计:")
        print(f"  成功: {results['success_count']} 条")
        print(f"  失败: {results['error_count']} 条")

        if results['error_count'] > 0:
            print(f"\n导入完成，但有 {results['error_count']} 个错误，请检查日志")
        else:
            print(f"\n导入完成，所有数据导入成功！")

        print("="*60)


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="简化版污染类型数据导入工具")
    parser.add_argument("excel_file", help="Excel文件路径")
    parser.add_argument("--customer-id", default="default", help="客户标识 (默认: default)")
    parser.add_argument("--clear-existing", action="store_true", help="清空现有数据")

    args = parser.parse_args()

    importer = SimplePollutionTypesImporter()

    try:
        importer.import_from_excel(args.excel_file, args.customer_id, args.clear_existing)
    except Exception as e:
        logger.error(f"导入失败: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()